import asyncio
from dotenv import load_dotenv
#import os
import logging
from datetime import datetime
from pathlib import Path
import nltk
from nltk.tokenize import word_tokenize
load_dotenv()

if __name__ == "__main__":
    import nest_asyncio
    nest_asyncio.apply()


    async def main():

        draft_iteration_output_dir = Path(r"draft_iteration_output")
        draft_iteration_output_dir.mkdir(exist_ok=True)

        final_survey_output_dir = Path(r"survey_output")
        final_survey_output_dir.mkdir(exist_ok=True)

        #topic = "What does the technology development roadmap for multi-modal large models look like?"
        #topic = "What are the methods to enhance the planning capability of large-scale models, and what are their respective strengths and weaknesses?"
        topic ="a comprehensive review of loss function"
        from research_agent.core.pipeline import Pipeline

        console_handler = logging.StreamHandler()  # 创建一个输出到控制台的 Handler
        console_handler.setLevel(logging.DEBUG)  # 将 Logger 的日志级别设置为 INFO
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        console_handler.setFormatter(formatter)  # 将格式器添加到 Handler 中
        logger = logging.getLogger("my_logger")
        logger.setLevel(logging.DEBUG)
        logger.addHandler(console_handler)  # 将 Handler 添加到 Logger 中

        pipeline = Pipeline(logger=logger,output_dir=draft_iteration_output_dir)
        survey_markdown = await pipeline.iteration(topic)
        # 在survey_markdown的文件头追加文件标题
        #punctuations =['#','%','&','*','|','\\','、',"\"",'<','>','?','/']
        #file_name = ' '.join([word for word in word_tokenize(topic) if word not in punctuations])

        with open('review.md', 'w', encoding='utf-8') as file:
            file.write('# '+topic+'\n\n')
            file.write(survey_markdown)
    asyncio.run(main())
